"""
Talks to an XMLRPC server running inside of an active IDA Pro instance,
in order to query it about the database.  Allows symbol resolution and
interactive debugging.
"""

from __future__ import annotations

import errno
import functools
import socket
import sys
import time
import traceback
import xmlrpc.client
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Tuple
from typing import TypeVar

import gdb
from pygments import highlight
from pygments.formatters import Terminal256Formatter
from pygments.lexers import CppLexer
from typing_extensions import Concatenate
from typing_extensions import ParamSpec
from typing_extensions import override

import pwndbg
import pwndbg.aglib.arch
import pwndbg.aglib.elf
import pwndbg.aglib.memory
import pwndbg.aglib.regs
import pwndbg.aglib.vmmap
import pwndbg.decorators
import pwndbg.integration
import pwndbg.lib.cache
import pwndbg.lib.funcparser
from pwndbg.color import message
from pwndbg.dbg import EventType
from pwndbg.lib.functions import Function

ida_rpc_host = pwndbg.config.add_param("ida-rpc-host", "127.0.0.1", "ida xmlrpc server address")
ida_rpc_port = pwndbg.config.add_param("ida-rpc-port", 43718, "ida xmlrpc server port")
ida_timeout = pwndbg.config.add_param("ida-timeout", 2, "time to wait for ida xmlrpc in seconds")


_ida: xmlrpc.client.ServerProxy | None = None

# to avoid printing the same exception multiple times, we store the last exception here
_ida_last_exception: BaseException | None = None

# to avoid checking the connection multiple times with no delay, we store the last time we checked it
_ida_last_connection_check = 0

P = ParamSpec("P")
T = TypeVar("T")


@pwndbg.decorators.only_after_first_prompt()
@pwndbg.config.trigger(ida_rpc_host, ida_rpc_port, pwndbg.integration.provider_name, ida_timeout)
def ida_config_changed() -> None:
    if pwndbg.integration.provider_name.value == "ida":
        # We need to (re)connect the client, possibly with updated values.
        try_init_ida_rpc_client()


def ensure_disabled() -> None:
    global _ida
    _ida = None
    pwndbg.integration.unset_provider()
    pwndbg.integration.provider_name.value = "none"


def try_init_ida_rpc_client() -> bool:
    """
    Try to connect to the IDA RPC client.

    If the connection succeeds, or we were already connected,
    return True. Otherwise, False.

    An appropriate message will be also printed to the user.
    """

    global _ida, _ida_last_exception, _ida_last_connection_check

    xmlrpc.client.MAXINT = 10**100  # type: ignore[misc]
    xmlrpc.client.MININT = -(10**100)  # type: ignore[misc]

    now = time.time()
    if _ida is None and (now - _ida_last_connection_check) < int(ida_timeout) + 5:
        return False

    addr = f"http://{ida_rpc_host}:{ida_rpc_port}"

    _ida = xmlrpc.client.ServerProxy(addr)
    socket.setdefaulttimeout(int(ida_timeout))

    exception = None  # (type, value, traceback)
    try:
        _ida.here()
        idc._update()
        pwndbg.integration.set_provider(IdaProvider())

        print(message.success(f"Pwndbg successfully connected to Ida Pro xmlrpc: {addr}"))

        if pwndbg.integration.provider_name.value != "ida":
            # We managed to successfully connect, and this happened because an Ida
            # command was invoked, rather than the user setting the integration-provider parameter.
            # So, we want to set the provider name now.
            # Note that ida_config_changed() is a trigger, and not a value listener, so it won't
            # be called when we set the value here (which is good, we would have recursion otherwise).
            pwndbg.integration.provider_name.value = "ida"

            assert (
                len(pwndbg.config.triggers[pwndbg.integration.provider_name.name]) == 2
            ), "Does this new function need to be called here?"

        return True

    except (TimeoutError, xmlrpc.client.ProtocolError):
        exception = sys.exc_info()
    except OSError as e:
        if e.errno == errno.ECONNREFUSED:
            print(
                message.error("Connection refused. ")
                + message.hint("Did you start ./ida_script.py from Ida?")
            )
        else:
            exception = sys.exc_info()

    ensure_disabled()

    if exception:
        if (
            not isinstance(_ida_last_exception, exception[0])
            or _ida_last_exception.args != exception[1].args
        ):
            if hasattr(pwndbg.config, "exception_verbose") and pwndbg.config.exception_verbose:
                print(message.error("[!] Ida Pro xmlrpc error"))
                traceback.print_exception(*exception)
            else:
                exc_type, exc_value, _ = exception
                print(
                    message.error(
                        f"Failed to connect to IDA Pro ({exc_type.__qualname__}: {exc_value})"
                    )
                )
                if exc_type is socket.timeout:
                    print(
                        message.notice("To increase the time to wait for IDA Pro use `")
                        + message.hint("set ida-timeout <new-timeout-in-seconds>")
                        + message.notice("`")
                    )
                else:
                    print(
                        message.notice("For more info invoke `")
                        + message.hint("set exception-verbose on")
                        + message.notice("`")
                    )
                print(
                    message.notice("To disable IDA Pro integration invoke `")
                    + message.hint("set integration-provider none")
                    + message.notice("`")
                )

    _ida_last_exception = exception and exception[1]
    _ida_last_connection_check = now
    return False


# We cannot catch the ConnectionRefusedError here, nor in @withIDA because there
# may be multiple nested decorated functions, and the bottom most one will swallow
# the exception up prevent it from bubbling to the top. Thus, we catch
# ConnectionRefusedError in CommandObj.__call__().
def enabledIDA(func: Callable[P, T]) -> Callable[P, T | None]:
    """
    If we have a connection to Ida, call the function.

    Otherwise, return None. Thus, all functions decorated with this must have
    "| None" in their return signature.

    This will not try to open a connection if it doesn't already exist.
    No messages will be printed.
    """

    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T | None:
        if _ida is None:
            assert pwndbg.integration.provider_name.value != "ida"
            return None

        return func(*args, **kwargs)

    return wrapper


def establish_connection() -> bool:
    """
    If we already had a connection, or succeed in creating a new one, return True.
    Otherwise False.
    """
    if _ida is not None:
        return True

    print(message.notice("Trying to connect to Ida..."))
    ok = try_init_ida_rpc_client()

    if not ok:
        print(message.error("Aborting."))
        return False

    return True


def withIDA(func: Callable[P, T]) -> Callable[P, T | None]:
    """
    Try to connect to Ida before running the decorated function.

    If we fail connecting to Ida, return None. Thus, all functions
    decorated with this must have "| None" in their return signature.

    Use this for user-initiated stuff like pwndbg.commands.ida.save_ida().
    """

    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T | None:
        if establish_connection():
            return func(*args, **kwargs)
        return None

    return wrapper


def withHexrays(func: Callable[P, T]) -> Callable[P, T | None]:
    @enabledIDA
    @functools.wraps(func)
    def wrapper(*a: P.args, **kw: P.kwargs) -> T | None:
        if _ida is not None and _ida.init_hexrays_plugin():
            return func(*a, **kw)
        return None

    return wrapper


def takes_address(function: Callable[Concatenate[int, P], T]) -> Callable[Concatenate[int, P], T]:
    @functools.wraps(function)
    def wrapper(address: int, *args: P.args, **kwargs: P.kwargs) -> T:
        return function(l2r(address), *args, **kwargs)

    return wrapper


def returns_address(function: Callable[P, int]) -> Callable[P, int]:
    @functools.wraps(function)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> int:
        return r2l(function(*args, **kwargs))

    return wrapper


def l2r(addr: int) -> int:
    region_start = pwndbg.aglib.vmmap.addr_region_start(addr)
    if region_start is None:
        return 0
    result = (addr - region_start + base()) & pwndbg.aglib.arch.ptrmask
    return result


def r2l(addr: int) -> int:
    region_start = pwndbg.aglib.vmmap.addr_region_start(addr)
    if region_start is None:
        return 0
    result = (addr - base() + region_start) & pwndbg.aglib.arch.ptrmask
    return result


def remote(function) -> None:
    """Runs the provided function in IDA's interpreter.

    The function must be self-contained and not reference any
    global variables."""


@pwndbg.lib.cache.cache_until("objfile")
def base():
    segaddr: int = _ida.get_next_seg(0)
    base: int = _ida.get_fileregion_offset(segaddr)

    return segaddr - base


@enabledIDA
@takes_address
def Comment(addr: int):
    return _ida.get_cmt(addr, 0) or _ida.get_cmt(addr)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def Name(addr: int):
    return _ida.get_name(addr, 0x1)  # GN_VISIBLE


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def GetFuncOffset(addr: int):
    rv = _ida.get_func_off_str(addr)
    return rv


@enabledIDA
@takes_address
def GetFuncAttr(addr: int, attr: int):
    rv = _ida.get_func_attr(addr, attr)
    return rv


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def GetType(addr: int):
    rv = _ida.get_type(addr)
    return rv


@enabledIDA
@returns_address
def here() -> int:
    return _ida.here()  # type: ignore[return-value]


@enabledIDA
@takes_address
def Jump(addr: int):
    # uses C++ api instead of idc one to avoid activating the IDA window
    return _ida.jumpto(addr, -1, 0)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def Anterior(addr: int):
    hexrays_prefix = b"\x01\x04; "
    lines = []
    for i in range(10):
        r: bytes | None = _ida.get_extra_cmt(addr, 0x3E8 + i)  # E_PREV
        if not r:
            break
        if r.startswith(hexrays_prefix):
            r = r[len(hexrays_prefix) :]
        lines.append(r)
    return b"\n".join(lines)


@enabledIDA
def GetBreakpoints():
    for i in range(GetBptQty()):
        yield GetBptEA(i)


@enabledIDA
def GetBptQty():
    return _ida.get_bpt_qty()


@enabledIDA
@returns_address
def GetBptEA(i: int) -> int:
    return _ida.get_bpt_ea(i)  # type: ignore[return-value]


_breakpoints: List[gdb.Breakpoint] = []


@pwndbg.dbg.event_handler(EventType.CONTINUE)
@pwndbg.dbg.event_handler(EventType.STOP)
@enabledIDA
def UpdateBreakpoints() -> None:
    # XXX: Remove breakpoints from IDA when the user removes them.
    current = {eval(b.location.lstrip("*")) for b in _breakpoints}
    want = set(GetBreakpoints())

    for addr in current - want:
        for bp in _breakpoints:
            if int(bp.location.lstrip("*"), 0) == addr:
                bp.delete()
                break
        _breakpoints.remove(bp)

    for addr in want - current:
        if not pwndbg.aglib.memory.peek(addr):
            continue

        bp = gdb.Breakpoint("*" + hex(int(addr)))
        _breakpoints.append(bp)


@enabledIDA
@takes_address
def SetColor(pc, color):
    return _ida.set_color(pc, 1, color)


colored_pc = None


@pwndbg.dbg.event_handler(EventType.STOP)
@enabledIDA
def Auto_Color_PC() -> None:
    global colored_pc
    colored_pc = pwndbg.aglib.regs.pc
    SetColor(colored_pc, 0x7F7FFF)


@pwndbg.dbg.event_handler(EventType.CONTINUE)
@enabledIDA
def Auto_UnColor_PC() -> None:
    global colored_pc
    if colored_pc:
        SetColor(colored_pc, 0xFFFFFF)
    colored_pc = None


@enabledIDA
@returns_address
@pwndbg.lib.cache.cache_until("objfile")
def LocByName(name) -> int:
    return _ida.get_name_ea_simple(str(name))  # type: ignore[return-value]


@enabledIDA
@takes_address
@returns_address
@pwndbg.lib.cache.cache_until("objfile")
def PrevHead(addr):
    return _ida.prev_head(addr)


@enabledIDA
@takes_address
@returns_address
@pwndbg.lib.cache.cache_until("objfile")
def NextHead(addr):
    return _ida.next_head(addr)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def GetFunctionName(addr):
    return _ida.get_func_name(addr)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def GetFlags(addr):
    return _ida.get_full_flags(addr)


@enabledIDA
@pwndbg.lib.cache.cache_until("objfile")
def isASCII(flags):
    return _ida.is_strlit(flags)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("objfile")
def ArgCount(address) -> None:
    pass


@enabledIDA
def SaveBase(path: str):
    return _ida.save_database(path)


@enabledIDA
def GetIdbPath():
    return _ida.get_idb_path()


@takes_address
@pwndbg.lib.cache.cache_until("stop")
def has_cached_cfunc(addr):
    return _ida.has_cached_cfunc(addr)


_lexer = CppLexer()
_formatter = Terminal256Formatter(style="monokai")


@withHexrays
@takes_address
@pwndbg.lib.cache.cache_until("stop")
def decompile(addr) -> str | None:
    code: str | None = _ida.decompile(addr)
    if code is not None and not pwndbg.config.disable_colors:
        code = highlight(code, _lexer, _formatter)
    return code


@withHexrays
@takes_address
@pwndbg.lib.cache.cache_until("stop")
def decompile_context(pc, context_lines) -> str | None:
    code: str | None = _ida.decompile_context(pc, context_lines)
    if code is not None and not pwndbg.config.disable_colors:
        code = highlight(code, _lexer, _formatter)
    return code


@enabledIDA
@pwndbg.lib.cache.cache_until("forever")
def get_ida_versions() -> Dict[str, str]:
    return _ida.versions()  # type: ignore[return-value]


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetStrucQty():
    return _ida.get_struc_qty()


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetStrucId(idx):
    return _ida.get_struc_by_idx(idx)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetStrucName(sid):
    return _ida.get_struc_name(sid)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetStrucSize(sid):
    return _ida.get_struc_size(sid)


@enabledIDA
@takes_address
@pwndbg.lib.cache.cache_until("stop")
def GetFrameId(addr):
    return _ida.get_frame_id(addr)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberQty(sid):
    return _ida.get_member_qty(sid)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberSize(sid, offset):
    return _ida.get_member_size(sid, offset)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberId(sid, offset):
    return _ida.get_member_id(sid, offset)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberName(sid, offset):
    return _ida.get_member_name(sid, offset)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberOffset(sid, member_name):
    return _ida.get_member_offset(sid, member_name)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetMemberFlag(sid, offset):
    return _ida.get_member_flag(sid, offset)


@enabledIDA
@pwndbg.lib.cache.cache_until("stop")
def GetStrucNextOff(sid, offset):
    return _ida.get_next_offset(sid, offset)


class IDC:
    query = "{k:v for k,v in globals()['idc'].__dict__.items() if isinstance(v, int)}"

    def _update(self) -> None:
        data: Dict[Any, Any] = _ida.eval(self.query)
        self.__dict__.update(data)


idc = IDC()


def print_member(sid, offset) -> None:
    mname = GetMemberName(sid, offset) or "(no name)"
    msize = GetMemberSize(sid, offset) or 0
    print(f"    +{offset:#x} - {mname} [{msize:#x} bytes]")


def print_structs() -> None:
    for i in range(GetStrucQty() or 0):
        sid = GetStrucId(i)

        name = GetStrucName(sid)
        size = GetStrucSize(sid)

        print(f"{name} - {size:#x} bytes")

        offset = 0
        while offset < size:
            print_member(sid, offset)
            offset = GetStrucNextOff(sid, offset)


ida_replacements = {
    "__int64": "signed long long int",
    "__int32": "signed int",
    "__int16": "signed short",
    "__int8": "signed char",
    "__uint64": "unsigned long long int",
    "__uint32": "unsigned int",
    "__uint16": "unsigned short",
    "__uint8": "unsigned char",
    "_BOOL_1": "unsigned char",
    "_BOOL_2": "unsigned short",
    "_BOOL_4": "unsigned int",
    "_BYTE": "unsigned char",
    "_WORD": "unsigned short",
    "_DWORD": "unsigned int",
    "_QWORD": "unsigned long long",
    "__pure": "",
    "__hidden": "",
    "__return_ptr": "",
    "__struct_ptr": "",
    "__array_ptr": "",
    "__fastcall": "",
    "__cdecl": "",
    "__thiscall": "",
    "__userpurge": "",
}


class IdaProvider(pwndbg.integration.IntegrationProvider):
    @pwndbg.decorators.suppress_errors()
    @enabledIDA
    def get_symbol(self, addr: int) -> str | None:
        exe = pwndbg.aglib.elf.exe()
        if exe:
            exe_map = pwndbg.aglib.vmmap.find(exe.address)
            if exe_map and addr in exe_map:
                return Name(addr) or GetFuncOffset(addr) or None
        return None

    @pwndbg.decorators.suppress_errors(fallback=())
    def get_versions(self) -> Tuple[str, ...]:
        ida_versions = get_ida_versions()

        if ida_versions is not None:
            ida_version = f"IDA PRO:  {ida_versions['ida']}"
            ida_py_ver = f"IDA Py:   {ida_versions['python']}"
            ida_hr_ver = f"Hexrays:  {ida_versions['hexrays']}"
            return (ida_version, ida_py_ver, ida_hr_ver)
        return ()

    @pwndbg.decorators.suppress_errors(fallback=True)
    @enabledIDA
    def is_in_function(self, addr: int) -> bool:
        return bool(GetFunctionName(addr))

    @pwndbg.decorators.suppress_errors(fallback=[])
    @enabledIDA
    def get_comment_lines(self, addr: int) -> List[str]:
        pre = Anterior(addr)
        return pre.decode().split("\n") if pre else []

    @pwndbg.decorators.suppress_errors()
    @enabledIDA
    def decompile(self, addr: int, lines: int) -> List[str] | None:
        code = decompile_context(addr, lines // 2)
        if code:
            return code.splitlines()
        else:
            return None

    @pwndbg.decorators.suppress_errors()
    @enabledIDA
    def get_func_type(self, addr: int) -> Function | None:
        typename: str = GetType(addr)

        if typename:
            typename += ";"

            # GetType() does not include the name.
            typename = typename.replace("(", " function_name(", 1)

            for k, v in ida_replacements.items():
                typename = typename.replace(k, v)

            return pwndbg.lib.funcparser.ExtractFuncDeclFromSource(typename + ";")

        return None

    @override
    def disable(self) -> None:
        ensure_disabled()
